import torch
import torch.nn as nn
from perceptron_model import Perceptron  # 导入单层感知机模型类
from net_model import MultiLayerPerceptron  # 导入多层感知机模型类

# 检查是否有可用的CUDA设备,如果有则使用GPU,否则使用CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"using device: {device}")

# 定义数据
X = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0]).view(-1, 1).to(device) # 定义输入数据X，改变结构为了之后前向传播使用
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0]).view(-1, 1).to(device)  # 定义标签数据y，改变结构为了之后前向传播使用

# 实例化单层感知机模型
perceptron_model = Perceptron().to(device)

# 定义损失函数和优化器
loss_fnc = nn.MSELoss()  # 定义损失函数为均方误差
optimizer = torch.optim.SGD(perceptron_model.parameters(), lr=0.01)  # 定义优化器为随机梯度下降, 学习率为0.01

# 训练单层感知机模型
for epoch in range(500):
    y_pred = perceptron_model(X)  # 使用模型对输入X进行预测，得到预测结果y_pred
    loss = loss_fnc(y_pred, y)  # 计算预测结果y_pred和真实值y之间的损失

    optimizer.zero_grad()  # 清零优化器中的梯度信息，以便于进行下一轮的优化
    loss.backward()  # 对损失进行反向传播，计算梯度
    optimizer.step()  # 使用优化器对模型参数进行更新

    if (epoch + 1) % 50 == 0:  # 每50次训练，打印一次当前的训练轮数和损失
        print(f'Epoch [{epoch + 1}/500], Loss: {loss.item():.4f}')

# 保存单层感知机模型
torch.save(perceptron_model.state_dict(), 'perceptron_model.pth')
print('Saved perceptron model to perceptron_model.pth')

# 实例化多层感知机模型
net_model = MultiLayerPerceptron().to(device)

# 定义损失函数和优化器
loss_fnc = nn.MSELoss()  # 定义损失函数为均方误差
optimizer = torch.optim.SGD(net_model.parameters(), lr=0.01)  # 定义优化器为随机梯度下降, 学习率为0.01

# 训练多层感知机模型
for epoch in range(500):
    y_pred = net_model(X)  # 使用模型对输入X进行预测，得到预测结果y_pred
    loss = loss_fnc(y_pred, y)  # 计算预测结果y_pred和真实值y之间的损失

    optimizer.zero_grad()  # 清零优化器中的梯度信息，以便于进行下一轮的优化
    loss.backward()  # 对损失进行反向传播，计算梯度
    optimizer.step()  # 使用优化器对模型参数进行更新

    if (epoch + 1) % 50 == 0:  # 每50次训练，打印一次当前的训练轮数和损失
        print(f'Epoch [{epoch + 1}/500], Loss: {loss.item():.4f}')

# 保存多层感知机模型
torch.save(net_model.state_dict(), 'net_model.pth')
print('Saved net model to net_model.pth')
